import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
h = 8.6125 # 空气交换系数
# 录入材料信息
# 根据厚度对介质进行分割，方便x轴标数计算
m = np.array([6, 60, 36, 50])
# 分割后计算总厚度，用于规定X的范围
m_sum = np.sum(m)
# 对时间进行分割和规定总时长
n = 3600
time = 3600
# 材料的厚度
# 四种材料厚度
l1 = 0.6/1000
l2 = 6/1000
l3 = 3.6/1000
l4 = 5/1000
l = np.array([l1, l2, l3, l4])
# 四种材料的热传导率
lam_1 = 0.082
lam_2 = 0.37
lam_3 = 0.045
lam_4 = 0.028
lam = np.array([lam_1, lam_2, lam_3, lam_4])
# 四种材料的密度
de_1 = 300
de_2 = 862
de_3 = 74.2
de_4 = 1.18
de = np.array([de_1, de_2, de_3, de_4])
# 四种材料的比热容
c1 = 1377
c2 = 2100
c3 = 1726
c4 = 1005
c = np.array([c1, c2, c3, c4])


# 计算热扩率 && 计算材料长度分割和时间步长分割求解 && 计算各层介质剖分的步长比
a = []  # 热扩率
derta_x = []  # 材料分割步长
derta_t = time/n  # 时间分割步长
r = []  # 介质剖分步长
u = np.zeros([m_sum, n])  # 定义四层耦合介质温度分布矩阵 每一列是一个x轴
for i in range(len(c)):

    temp_a = lam[i]/(c[i]*de[i])
    temp_x = l[i] / m[i]
    temp_r = derta_t/np.square(temp_x)*temp_a

    a.append(temp_a)
    derta_x.append(temp_x)
    r.append(temp_r)
a = np.array(a)
derta_x = np.array(derta_x)
r = np.array(r)
# print(a,derta_x,r)

# 初始热量和恒定边界热量
u[:, 0] = 37
u[0, :] = 75
# print(u)

# 构造差分格式的系数矩阵
A = np.zeros([m_sum, m_sum])
# 第一个边界
for i in range(0, m[0]-1):
    A[i, i] = 1+2*r[0]
    A[i, i+1] = -r[0]
    if i >= 1:
        A[i, i-1] = -r[0]
# 第一个边界的边界条件
A[m[0]-1, m[0]-1
] = (lam[0]/derta_x[0]+lam[1]/derta_x[1])
A[m[0]-1, m[0]-2] = lam[0]/derta_x[0]
A[m[0]-1, m[0]] = lam[1]/derta_x[1]

# 第二个边界
for i in range(m[0], m[1]+m[0]-1):
    A[i, i] = 1 + 2 * r[1]
    A[i, i + 1] = -r[1]
    A[i, i - 1] = -r[1]
# 第二个边界的边界条件
A[m[1]+m[0]-1, m[1]+m[0]-1] = (lam[2]/derta_x[2]+lam[1]/derta_x[1])
A[m[1]+m[0]-1, m[1]+m[0]-2] = lam[1]/derta_x[1]
A[m[1]+m[0]-1, m[1]+m[0]] = lam[2]/derta_x[2]

# 第三个边界
for i in range(m[1]+m[0], m[1]+m[0]+m[2]-1):
    A[i, i] = 1 + 2 * r[2]
    A[i, i + 1] = -r[2]
    A[i, i - 1] = -r[2]
# 第三个边界的边界条件
A[m[2]+m[1]+m[0]-1, m[2]+m[1]+m[0]-1] = (lam[3]/derta_x[3]+lam[2]/derta_x[2])
A[m[2]+m[1]+m[0]-1, m[2]+m[1]+m[0]-2] = lam[2]/derta_x[2]
A[m[2]+m[1]+m[0]-1, m[2]+m[1]+m[0]] = lam[3]/derta_x[3]

# 第四个边界
for i in range(m[2]+m[1]+m[0], m[3]+m[1]+m[0]+m[2]-1):
    A[i, i] = 1 + 2 * r[3]
    A[i, i + 1] = -r[3]
    A[i, i - 1] = -r[3]
A[m_sum-1, m_sum-1] = h + lam[3]/derta_x[3]
A[m_sum-1, m_sum-2] = lam[3]/derta_x[3]

# 构造右端项
for k in range(1, n):
    b = np.zeros([m_sum, 1])  # 临时列变量
    for i in range(1, m_sum-1):
        b[i, 0] = u[i, k-1]  # 固定时间，遍历x轴
    # 处理当前列向量存在的边界条件
    b[0] = u[1, k-1] + r[0]*u[0, k]  # 两个恒温源,当前时间起始处的温度始终受到37度和75度的影响
    b[m[0]-1, 0] = 0
    b[m[0]+m[1]-1, 0] = 0
    b[m[0]+m[1]+m[2]-1, 0] = 0
    b[m_sum-1, 0] = 37*h
    # print(b)
    # 差分追赶法
    bb = np.diag(A).T  # 获得A对角线上的值
    aa = np.diag(A, -1).T  # 获取对角线旁边对角的值
    aa = np.insert(aa, 0, 0) # 因为长度和计算原因，在前面补0
    c = np.diag(A, 1).T  # 获取对角线旁边对角的值
    c = np.append(c,0)
    N = len(bb)  # 获取对角线上值的长度
    L = np.zeros(N)  # 创建一个长度和对角线元素相等的零矩阵
    y = np.zeros(N)
    x = np.zeros(N)
    uu = np.zeros(N)
    # 一些初始值
    uu0 = 0
    y0 = 0
    aa[0] = 0
    L[0] = bb[0]-aa[0]*uu0
    y[0] = (b[0] - y0*aa[0])/L[0]
    uu[0] = c[0]/L[0]
    for i in range(1, N):
        L[i] = bb[i] - aa[i] * uu[i - 1]
        uu[i-1] = c[i-1] / L[i]
        y[i] = (b[i] - y[i - 1] * aa[i]) / L[i]
    L[N - 1] = bb[N - 1] - aa[N - 1] * uu[N - 2]
    y[N - 1] = (b[N - 1] - y[N - 2] * aa[N - 1]) / L[N - 1]
    x[N - 1] = y[N - 1]

    for i in range(N - 2, -1, -1):
        x[i] = y[i] - uu[i] * x[i + 1]
    # print(x.shape)
    u[1:m_sum+1, k] = np.array(x[1:]).T
    '''
        
    
    '''
    # print(y)

# print(u)
# 画图
X_x = np.array(range(0,m_sum))
fig = plt.figure("3D Surface", facecolor="lightgray")
plt.title("Temperature3D", fontsize=18)

# 设置为3D图片类型
ax3d = Axes3D(fig, auto_add_to_figure=False)
fig.add_axes(ax3d)

ax3d.set_xlabel("time")
ax3d.set_ylabel("x")
ax3d.set_zlabel("temperature")
plt.tick_params(labelsize=10)
t = np.array(range(0,time))
x_x, y_y = np.meshgrid(t, X_x)
ax3d.plot_surface(x_x, y_y, u, cstride=20, rstride=20, cmap="jet")

# plt.savefig("Temperature3D.png")  # 保存图片
plt.show()